{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "None\n",
      "None\n",
      "None\n",
      "None\n",
      "None\n",
      "None\n"
     ]
    }
   ],
   "source": [
    "import cv2\n",
    "import os\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "category_dict = {\"red\": 1,\"green\": 2,\"yellow\": 3,\"red_left\": 4,\"red_right\": 5,\"yellow_left\": 6, \"yellow_right\": 7,\"green_left\": 8,\n",
    "                 \"green_right\": 9,\"red_forward\": 10,\"green_forward\": 11,\"yellow_forward\": 12,\"horizon_red\": 13,\"horizon_green\": 14,\n",
    "                 \"horizon_yellow\": 15,\"off\": 16,\"traffic_sign\": 17,\"car\": 18,\"motor\": 19,\"bike\": 20,\"bus\": 21,\"truck\": 22, \"suv\":23,\n",
    "                 \"express\": 24,\"person\":25}\n",
    "\n",
    "root_path = 'C:/Users/ZZK/Documents/auto_drive/train_dataset/'\n",
    "target_path = 'C:/Users/ZZK/Desktop/'\n",
    "with open(\"C:/Users/ZZK/Documents/auto_drive/train_dataset/train.txt\",\"r\") as outfile:\n",
    "    for line in outfile.readlines():\n",
    "        line = line.split(' ')\n",
    "        img_name = line[0]\n",
    "        line = line[2:]\n",
    "        bbox_num = len(line)\n",
    "        img_cv2 = cv2.imread(root_path + img_name)\n",
    "        for i in range(bbox_num):\n",
    "            line_num = line[i].split(',')\n",
    "            if len(line_num)<4:\n",
    "                break\n",
    "            x1 = int(line_num[0])\n",
    "            y1 = int(line_num[1])\n",
    "            x2 = int(line_num[2])\n",
    "            y2 = int(line_num[3])\n",
    "            value = int(line_num[4])\n",
    "            confidence = line_num[5]\n",
    "            text = list(category_dict.keys())[list(category_dict.values()).index(value)] + '  ' +confidence\n",
    "            img_cv2 = cv2.putText(img_cv2,text,(x1,y2-np.random.randint(-20,20)),cv2.FONT_HERSHEY_COMPLEX,0.3,(255,50,200),1)\n",
    "            cv2.line(img_cv2,(x1,y1),(x1,y2),(0,0,255),1,8)\n",
    "            cv2.line(img_cv2,(x1,y2),(x2,y2),(0,0,255),1,8)\n",
    "            cv2.line(img_cv2,(x1,y1),(x2,y1),(0,0,255),1,8)\n",
    "            cv2.line(img_cv2,(x2,y1),(x2,y2),(0,0,255),1,8)\n",
    "        cv2.imwrite(target_path + img_name,img_cv2)\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_to_write = open(\"C:/Users/ZZK/Desktop/repeat_box.txt\", 'w')#新生成的文件\n",
    "imglb = []\n",
    "corlb = []\n",
    "with open(\"C:/Users/ZZK/Desktop/1.3.txt\",\"r\") as outfile:\n",
    "    for line in outfile.readlines():\n",
    "        line = line.split(' ')\n",
    "        img_name = line[0]\n",
    "        line = line[2:]\n",
    "        bbox_num = len(line)\n",
    "\n",
    "        box = []\n",
    "        for i in range(bbox_num):\n",
    "            line_num = line[i].split(',')\n",
    "            if len(line_num)<4:\n",
    "                break\n",
    "            x1 = int(line_num[0])\n",
    "            y1 = int(line_num[1])\n",
    "            x2 = int(line_num[2])\n",
    "            y2 = int(line_num[3])\n",
    "            box += [line_num[:4]]\n",
    "            conf = int(line_num[5])\n",
    "\n",
    "        new_set = box\n",
    "        if bbox_num>1:\n",
    "            for i,item in enumerate (box):\n",
    "                if new_set.count(item) > 1:\n",
    "                    imglb += [img_name]\n",
    "                    corlb += [i]\n",
    "                    print(img_name,item)\n",
    "                    \n",
    "imglab = []\n",
    "corlab = []\n",
    "for j in range(len(imglb)):\n",
    "    if (j+1)%2 == 0:\n",
    "        imglab += [imglb[j]]\n",
    "        corlab += [corlb[j-1]]\n",
    "# print(imglab)\n",
    "# print(corlab)\n",
    "k = 0\n",
    "with open(\"C:/Users/ZZK/Desktop/1.3.txt\",\"r\") as outfile:\n",
    "    for lines in outfile.readlines():\n",
    "        line = lines.split(' ')\n",
    "        img_name = line[0]\n",
    "        img_name1 = line[1]\n",
    "        line = line[2:]\n",
    "        if img_name == imglab[k]:\n",
    "            file_to_write.write(img_name+' '+img_name1)\n",
    "            for l in range(len(line)):\n",
    "                if l != corlab[k]:\n",
    "                    file_to_write.write(' '+line[l])    \n",
    "            if imglab[k] == imglab[k+1]:\n",
    "                k += 1\n",
    "                if imglab[k] == imglab[k+1]:\n",
    "                    k += 1\n",
    "            k+=1\n",
    "            \n",
    "        else:\n",
    "            file_to_write.write(lines)\n",
    "        \n",
    "file_to_write.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
